from geesibling.core.lib._graph import DataType,DeviceType,Device


__doc__ = """
determine the number of gpus to use
Author: jirongting
datetime: 2024.04.15
version: 1 2024.04.15 first commit
"""

def device_config(attrs):
    d = []
    for k, v in attrs.items():
        d.append(Device(v["type"], k, v["memory"], v["free_memory"], v["execute_time"]))
    return d

"""
Parameters:
    ngpu - the number of gpus to use
"""
def device_num(ngpu):
    device_dist={}
    gpu_info={
            "type": DeviceType.gpu,
            "memory": 512 * 1024 * 1024 * 1024,
            "free_memory": 512 * 1024 * 1024 * 1024,
            "execute_time": 0,
    }
    if isinstance(ngpu, int):  # 检查 ngpu 是否为整数
        gpu_indices = range(ngpu)  # 如果是整数，创建一个从0到ngpu-1的序列
    elif isinstance(ngpu, list):  # 检查 ngpu 是否为列表
        gpu_indices = ngpu  # 如果是列表，直接使用列表
    else:
        raise ValueError("Invalid input: ngpu must be either an int or a list of indices.")
    
    
    for index in gpu_indices:
        key = "gpu:" + str(index)
        device_dist[key]=gpu_info
    
    return device_config(device_dist)
